"""
ocr_pipeline.py
一站式中文/英文混合 OCR 识别流水线
整套代码 = “诊断 + 预处理 + 多引擎择优 + 后处理 + 资源清理” 的端到端 OCR 工具箱，
既能在命令行一键跑，也能被其他 Python 项目 import后直接调用。
"""

import os
import re
import sys
import cv2
import numpy as np
import pytesseract
import tempfile
import functools
from PIL import Image
from typing import Tuple, Optional

# -------------------- 全局配置 --------------------
# 1. Windows 需手动指定 tesseract.exe 路径
TESSERACT_CMD = r"D:\programFile\Tesseract\tesseract.exe"
if os.name == "nt":
    pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD

# 2. EasyOCR Reader 单例
@functools.lru_cache(maxsize=1)
def _get_easyocr_reader():
    import easyocr
    return easyocr.Reader(
        ["ch_sim", "en"],
        gpu=False,
        model_storage_directory="./easyocr_models",
    )

# -------------------- 工具函数 --------------------
def print_image_info(image_path: str) -> None:
    """打印图像诊断信息"""
    print("\n" + "=" * 40)
    print("图像诊断信息".center(40))
    print("=" * 40)

    # 文件检查
    print(f"\n▶ 文件检查")
    print(f"路径: {image_path}")
    print(f"存在: {'是' if os.path.exists(image_path) else '否'}")
    if os.path.exists(image_path):
        print(f"大小: {os.path.getsize(image_path) / 1024:.1f} KB")
        try:
            with Image.open(image_path) as img:
                print(f"格式: {img.format}")
                print(f"原始尺寸: {img.size}")
        except Exception as e:
            print(f"Pillow读取失败: {e}")

    # OpenCV 检查
    print("\n▶ OpenCV检查")
    img_cv = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    if img_cv is None:
        print("❌ 读取失败 (可能路径编码问题)")
    else:
        print(f"尺寸: {img_cv.shape} (H×W×C)")
        print(f"数据类型: {img_cv.dtype}")
        print(f"亮度范围: {img_cv.min()}~{img_cv.max()}")

"""图像预处理（enhanced_preprocess）
• 做了什么：把原始图片变成“干净、高对比、无噪点”的二值图，再存成临时 jpg。
• 为什么要做：Tesseract/EasyOCR 对光照、模糊、背景噪点非常敏感，预处理能把识别率从 40% 提到 90% 以上。
• 怎么做（三步）：
① CLAHE 局部直方图均衡 → 解决光照不均；
② 自适应阈值（Gaussian）→ 动态二值化；
③ 形态学闭运算 → 去掉小黑点、连接断裂笔画。
• 细节：
– 临时文件用 tempfile.mkstemp，防止只读目录报错；
– 窗口调试仅在 SHOW_DEBUG_WINDOW=1 时弹出。
"""
def enhanced_preprocess(image_path: str) -> str:
    """增强版预处理，返回临时文件路径"""
    img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("无法读取图像文件")

    # 1. CLAHE 对比度增强
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    limg = cv2.merge([clahe.apply(l), a, b])
    enhanced = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

    # 2. 自适应二值化
    gray = cv2.cvtColor(enhanced, cv2.COLOR_BGR2GRAY)
    img_h, img_w = gray.shape
    block_size = max(31, min(img_w // 10, 101))
    thresh = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        block_size,
        8,
    )

    # 3. 形态学降噪
    kernel_size = max(1, min(img_w // 300, 3))
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
    cleaned = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

    # 保存到临时文件
    temp_fd, temp_path = tempfile.mkstemp(suffix=".jpg")
    os.close(temp_fd)
    cv2.imencode(".jpg", cleaned)[1].tofile(temp_path)

    # 可选：显示调试窗口
    if __debug__ and os.getenv("SHOW_DEBUG_WINDOW") == "1":
        cv2.imshow("Original", cv2.resize(img, (800, 600)))
        cv2.imshow("Processed", cv2.resize(cleaned, (800, 600)))
        cv2.waitKey(2000)
        cv2.destroyAllWindows()

    return temp_path

# -------------------- OCR 核心 --------------------
"""主 OCR（optimize_ocr）
• 做了什么：用 Tesseract 对预处理后的图做文字识别，并自动在多种 PSM 里挑最好的一组结果。
• 为什么要做：不同版式（单行、多行、竖排、稀疏文本）需要不同 PSM；挨个试比人工猜更稳健。
• 怎么做：
– 固定 oem 3（LSTM）+ 语言包 chi_sim+eng；
– 依次尝试 --psm 6/11/8/3，取最长结果；
– 过滤掉 ®©™° 等垃圾符号；
– 把结果送进 postprocess_text 做后处理。"""
def optimize_ocr(image_path: str) -> str:
    """Tesseract 多 PSM 择优"""
    base_config = (
        "--oem 3 "
        "-c preserve_interword_spaces=1 "
        '-c tessedit_char_blacklist=®©™°'
    )

    best_text = ""
    for psm in (6, 11, 8, 3):
        config = f"--psm {psm} " + base_config
        try:
            txt = pytesseract.image_to_string(
                image_path, lang="chi_sim+eng", config=config
            )
            if len(txt.strip()) > len(best_text.strip()):
                best_text = txt
        except Exception as e:
            print(f"Tesseract (psm={psm}) 报错: {e}")
            continue
    return postprocess_text(best_text)

"""备用 OCR（alternative_ocr）
• 做了什么：当 Tesseract 返回文本过短（<10 字符）时，自动切到 EasyOCR。
• 为什么要做：Tesseract 在复杂背景、艺术字、手写体上容易翻车，EasyOCR 的 CNN 模型能互补。
• 怎么做：
– 用 functools.lru_cache 把 Reader 做成单例，避免重复下载/加载模型；
– detail=0 直接拿纯文本；
– 结果同样走 postprocess_text。
"""
def alternative_ocr(image_path: str) -> str:
    """EasyOCR 备用方案"""
    try:
        reader = _get_easyocr_reader()
        results = reader.readtext(image_path, detail=0)
        return postprocess_text("\n".join(results))
    except ImportError:
        print("未安装EasyOCR，请运行: pip install easyocr")
        return ""
    except Exception as e:
        print(f"EasyOCR 错误: {e}")
        return ""

def postprocess_text(text: str) -> str:
    """后处理：空格/标点/孤立字符"""
    # 仅删除中英文之间多余的空格
    text = re.sub(r"([\u4e00-\u9fa5])\s+([a-zA-Z])", r"\1\2", text, flags=re.I)
    text = re.sub(r"([a-zA-Z])\s+([\u4e00-\u9fa5])", r"\1\2", text, flags=re.I)

    # 修正常见标点
    punctuation_map = {
        " ,": "，",
        " .": "。",
        " !": "！",
        " ?": "？",
        " :": "：",
        " ;": "；",
    }
    for wrong, right in punctuation_map.items():
        text = text.replace(wrong, right)

    # 去除孤立字符/空行
    lines = [ln.strip() for ln in text.splitlines() if len(ln.strip()) > 1]
    return "\n".join(lines)

# -------------------- 主流程 --------------------
def ocr_pipeline(image_path: str) -> Tuple[str, Optional[str]]:
    """完整 OCR 流水线"""
    print_image_info(image_path)

    try:
        processed_path = enhanced_preprocess(image_path)
        print(f"\n预处理完成 → {processed_path}")

        main_text = optimize_ocr(processed_path)
        print(f"\n主识别结果 ({len(main_text)}字符):\n{main_text[:200]}...")

        if len(main_text.strip()) < 10:
            alt_text = alternative_ocr(processed_path)
            print(f"\n备用识别结果 ({len(alt_text)}字符):\n{alt_text[:200]}...")
            final_text = alt_text if len(alt_text) > len(main_text) else main_text
            engine = "EasyOCR" if final_text == alt_text else "Tesseract"
        else:
            final_text, engine = main_text, "Tesseract"

        # 清理临时文件
        os.remove(processed_path)
        return final_text, engine

    except Exception as e:
        print(f"\n❌ 处理失败: {e}")
        return "", None

# -------------------- CLI --------------------
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("用法: python 65_一站式中英文混合OCR识别流水线.py <图片路径>")
        sys.exit(1)
    img_path = sys.argv[1]
    result, used_engine = ocr_pipeline(img_path)
    print("\n" + "=" * 40)
    print(f"最终结果 (使用 {used_engine})".center(40))
    print("=" * 40)
    print(result)

# python ./65_一站式中英文混合OCR识别流水线.py "D:\ideaSpace\rag-in-action-master\90-文档-Data\黑悟空\黑悟空英文.jpg"